"""
@Description :   2D heat transfer
@Author      :   Darcy 
@Time        :   2024/12/25 19:30:26
"""

import numpy as np
import time
from precision_data import fp  # floating point
from structured_mesh import *
from collocated_segregated import *


# clock
clock_begin = time.perf_counter() # define clock_begin

# main codes
dim = 2 # dimensions
ncx = 20 # number of cells in x direction
ncy = 20
ncz = 1

xmin = fp(0.0)
xmax = fp(0.883)

ymin = fp(0.0)
ymax = fp(0.83)

zmin = fp(0.0)
zmax = fp(1.0)

case = StructuredMesh()
case.createMesh(dim, ncx, ncy, ncz)
case.createCoordinates(xmin, xmax, ymin, ymax, zmin, zmax)
case.createFieldMeshData()

Tref = fp(373.0)
case.setInitialT(Tref)

case.createCoeffMeshData()

case.createSimulationData()
nsteps = 200
case.set_nsteps(nsteps)
dt = fp(1.0)
case.set_dt(dt)

case.createSolvingMethodData()
niter_T = 100
relax_T = fp(0.75)
res_T = fp(0.01)
temp_tol = fp(1.e-6)
case.set_temp_solver_param(niter_T, relax_T, res_T, temp_tol)

fluid = Fluid(ncx, ncy, ncz)
dens = fp(1.0)
mu = fp(0.001)
fluid.setInitialDenMu(dens, mu)

con = fp(81.0)
spht = fp(1.0)
fluid.set_con_spht(con, spht)

### 处理边界问题 1.46
fluidboundary = FluidBoundary(dim)

# 对某个cell的每一个faces, 设置其boundary
fluidboundary.createBoundaruOfCellFaces(case)

fluidboundary.createBoundaryData(dim,
                                 'wall', # xmin
                                 'wall', # xmax
                                 'wall', # ymin
                                 'wall', # ymax
                                 )

fluidboundary.createBoundaryDataTemp(dim,
                                     'wall', 'constant', fp(373.0), # xmin
                                     'wall', 'constant', fp(373.0), # xmax
                                     'wall', 'constant', fp(373.0), # ymin
                                     'wall', 'constant', fp(293.0), # ymax
                                     )


# Write VTK file
post = PostProcessing()
post.writeVTKCollocated_temp(case)

res_freq = 1
out_freq = 1000
post.Set_res_out_freq(res_freq, out_freq)


# Post-process for residual
nonlinsol_fname = post.nonlinsol_fname
linsol_fname = post.linsol_fname

# Open mass residual file
nonlinsol_fid = open(nonlinsol_fname, 'w')
nonlinsol_fid.write('#it, time, l2_T/l2_max_T\n')
nonlinsol_fid.close()

linsol_fid = open(linsol_fname, 'w')
linsol_fid.write('#it_nl, it, tot_it, norm, init, max, rel\n')
linsol_fid.close()

# Outer loop !!!!!!!!!!!!!!!!!!!
for it in range(1, nsteps+1):
    # Write control !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    if (it % 2 == 0 or it == 1 or it == nsteps):
        print('')
        print('--------------------')
        print('Begin iter = ', it)
        print('--------------------')
        
    case.stop_sim = collocated_segregated(it, case, fluid, fluidboundary, post)

    if (it % out_freq == 0 or it == nsteps or case.stop_sim):
        post.writeVTKCollocated_temp(case)

    if (it == 1 or it % res_freq == 0 or it == nsteps or case.stop_sim):
        print("it, time, l2_T/l2_max_T ", it, time.perf_counter() - clock_begin, case.l2_T/case.l2_max_T)

        with open(nonlinsol_fname, 'a') as nonlinsol_fid:
            nonlinsol_fid.write("{} {} {}\n".format(it, time.perf_counter() - clock_begin, case.l2_T/case.l2_max_T))

    # 判断是否满足迭代精度，是则退出迭代
    if case.stop_sim:
        break    


elapsed_time = time.perf_counter() - clock_begin
print('elapsed time: ', elapsed_time)